'''
Total reward:  576.0
mean reward:  5.76
'''

import gymnasium as gym
import ale_py
from lib import common

gym.register_envs(ale_py)

# Initialize the environment
# env = gym.make('ALE/FishingDerby-v5', render_mode='human')
env = gym.make('ALE/FishingDerby-v5')
# env = common.ProcessFrame84(gym.make('ALE/FishingDerby-v5'))
# env = gym.make('ALE/FishingDerby-v5, render_mode='rgb_array', frameskip=4, repeat_action_probability=0.0)
print("max max_episode_steps: ", env.spec.max_episode_steps)
count_frame = 0

# Reset the environment to get the initial state
total_reward = 0
# Run a loop to play the game
episoid = 100
for _ in range(episoid):
    state = env.reset()
    while True:
        # Take a random action
        # env.render()
        action = env.action_space.sample()

        # Get the next state, reward, done flag, and info from the environment
        state, reward, done, trunc, info = env.step(action)
        if reward > 0:
            total_reward += reward
            print("action: ", action)   
            print("reward: ", reward)
            print("info: ", info)

        # If done, reset the environment
        if done or trunc:
            print(f"done:{done} trunc:{trunc}")
            print("info: ", info)
            print("count_frame: ", count_frame)
            break

print("Total reward: ", total_reward)
print("mean reward: ", total_reward / episoid)

# Close the environment
env.close()